import matplotlib.pyplot as plt
import numpy as np

'''------------------------------------------- Residual-based adaptive refinement ---------------------------------------------'''
def get_top_k_indices(arr, k=100):
    '''get indices of the top k values in arr'''
    arr = arr.reshape(-1)  # reshape the array to 1D
    indices = np.argsort(arr)[::-1]  # sort indices in descending order
    return indices[:k]

def plot_truck_add_points_3d(points):
    fig = plt.figure(dpi=200)
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(points[:, 0], points[:, 1], points[:, 2], s=0.5)
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')  
    # plt.axis("scaled")
    plt.show()

def plot_points_3d(points):
    fig = plt.figure(dpi=200)
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(points[:, 0], points[:, 1], points[:, 2], s=0.5)
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')  
    # plt.axis("scaled")
    plt.show()

def convert_and_remove_duplicates(array):
    # Convert array of int32 to list
    array_list = [list(arr) for arr in array]
    # Flatten the 2D list
    flattened_list = [item for sublist in array_list for item in sublist]
    # Remove duplicate elements
    unique_list = list(set(flattened_list))
    return unique_list

def plot_points_xy(points, color="k", marker="."):
    figure = plt.figure()
    axis = figure.add_subplot(111)
    axis.scatter(points[:, 0], points[:, 1], s=0.1, color=color, marker=marker)
    # plt.axis("scaled")  # 设置x轴和y轴相同的缩放比例
    plt.show()

def tetrahedron_centroid(tetrahedrons, vertices):
    centroids = []
    for tetrahedron in tetrahedrons:
        vertex1 = vertices[tetrahedron[0]]
        vertex2 = vertices[tetrahedron[1]]
        vertex3 = vertices[tetrahedron[2]]
        vertex4 = vertices[tetrahedron[3]]
        centroid = (vertex1 + vertex2 + vertex3 + vertex4) / 4
        centroids.append(centroid)
    centroids_points = np.array(centroids)    
    return centroids_points

def tetrahedron_midpoint(tetrahedrons, vertices):
    midpoints = []
    for tetrahedron in tetrahedrons:
        vertex1 = vertices[tetrahedron[0]]
        vertex2 = vertices[tetrahedron[1]]
        vertex3 = vertices[tetrahedron[2]]
        vertex4 = vertices[tetrahedron[3]]
        midpoint12 = (vertex1 + vertex2) / 2
        midpoints.append(midpoint12)
        midpoint13 = (vertex1 + vertex3) / 2
        midpoints.append(midpoint13)
        midpoint14 = (vertex1 + vertex4) / 2
        midpoints.append(midpoint14)
        midpoint23 = (vertex2 + vertex3) / 2
        midpoints.append(midpoint23)
        midpoint24 = (vertex2 + vertex4) / 2
        midpoints.append(midpoint24)
        midpoint34 = (vertex3 + vertex4) / 2
        midpoints.append(midpoint34)
    midpoints = np.array(midpoints)    
    return midpoints

def remove_duplicate_points(points):
    unique_points = np.unique(points, axis=0)
    return unique_points

def remove_duplicate_rows(arr):
    sorted_rows = [tuple(sorted(row)) for row in arr]
    unique_rows = np.unique(sorted_rows, axis=0)
    return np.array(unique_rows)

def find_left_points(arr1, *arrs):
    # 创建布尔索引
    mask = np.ones(len(arr1), dtype=bool)
    for i, row in enumerate(arr1):
        for compare_row in np.concatenate(arrs):
            if np.array_equal(row, compare_row):
                mask[i] = False
                break
    # 使用布尔索引筛选出第一个数组中没有出现的元素
    result = arr1[mask]
    return result

def find_points_on_truck(all_mesh_points):
    x_min, x_max = 0, 4
    y_min, y_max = 80, 90
    z_min, z_max = 18.5, 20   
    boundary_points = []
    for point in all_mesh_points:
        x, y, z = point
        # 检查是否在长方体的边界上
        if np.isclose(x, x_min) or np.isclose(x, x_max):
            if y_min <= y <= y_max and z_min <= z <= z_max:
                boundary_points.append(point)
        elif np.isclose(y, y_min) or np.isclose(y, y_max):
            if x_min <= x <= x_max and z_min <= z <= z_max:
                boundary_points.append(point)
        elif np.isclose(z, z_min) or np.isclose(z, z_max):
            if x_min <= x <= x_max and y_min <= y <= y_max:
                boundary_points.append(point)
    return np.array(boundary_points)

def convert_small_numbers_to_zero(arr, threshold = 1e-8):
    condition = np.abs(arr) < threshold
    arr[condition] = 0
    return arr

    